#!/usr/bin/python3
#coding=utf8

import os,sys,glob,time
from threading import Thread
from imp import find_module,load_module,acquire_lock,release_lock

class PluginsManage(object):
    NOT_IMPORT = ['__init__.py']
    PLUGINS  = set()  #插件文件
    PAYLOADS = {}     #漏洞代码，一个插件里面可能包含数个漏洞。
    def __init__(self, userdir ='./',postfix='*.py'):
        self.userdir = userdir
        self.postfix = postfix

    def filter(self):
        return [p.split(os.sep)[-1] for p in glob.glob(self.userdir+'/'+self.postfix)]

    def modified(self):
        print(u"开始监控插件文件夹: %s"%self.userdir)
        before = dict([(f, None) for f in self.filter()])
        while True:
            time.sleep(10)
            after = dict([(f, None) for f in self.filter()])
            added = [f for f in after if not f in before]
            removed = [f for f in before if not f in after]
            if added:
                print(u"新增插件: %s"%added)
                self.add_plugins(added)
            if removed:
                print(u"删除插件: %s"%removed)
                self.del_plugins(removed)
            before = after

    def load_plugins(self,auto=False):
        plug = [plug for plug in self.filter()]
        self.add_plugins(plug)
        if auto:
            Thread(target=self.modified).start()

    def add_plugins(self,plugins):
        for plugin in plugins:
            if plugin not in self.NOT_IMPORT:
                fn = None
                name = plugin[:-3]  #.py是三位，如果postfix改了的话这里也要改。
                try:
                    acquire_lock()
                    fn, filename, desc = find_module(name, [self.userdir])
                    exec(fn.read())
                    #load_module(name, fn, filename, desc)
                    self.PLUGINS.add(plugin)
                except Exception as e:
                    print(u"插件: %s 导入出错,信息: %s"%(plugin,e))
                finally:
                    if fn:fn.close()
                    release_lock()

    def del_plugins(self,plugins):
        for plugin in plugins:
            if plugin in self.PLUGINS:
                self.PLUGINS.remove(plugin)

    @classmethod
    def register(cls,plug_name=None):
        def wrapper(plug):
            name = plug_name if plug_name else plug.__name__
            cls.PAYLOADS.update({name:plug()})
        return wrapper
#########################################################
class BasePlugin(object):
    #log = Logger()
    vul_name = ''
    vul_rank = ''
    vul_desc = ''
    vul_plan = ''
    def __init__(self):
        self.host = ''
        self.port = ''
        self.uri = ''
        self.payload = ''

class BaseHostPlugin(BasePlugin):
    service  = ''
    '''针对主机ip的插件，如redis未授权访问，'''
    def load(self,host,port):
        self.host = host
        self.port = port

class BaseHttpPlugin(BasePlugin):
    '''针对uri 如weblogic，tomcat等'''
    def load(self,uri):
        self.uri = uri

class BaseUrlPlugin(BasePlugin):
    '''针对url的，如sql注入，xss等'''
    pass

#########################################################
if __name__=='__main__':
    plug = PluginsManage('./')
    plug.load_plugins()
    for k,v in plug.PAYLOADS.items():
        print(v.vul_name)

